from kafka import KafkaProducer
from kafka.admin import KafkaAdminClient
from kafka.errors import KafkaError
import json
import logging
import time


class KafkaMessageProducer:
    def __init__(self, bootstrap_servers=['localhost:9092'], **kwargs):
        """
        初始化Kafka生产者

        Args:
            bootstrap_servers: Kafka服务器列表
            **kwargs: 其他Kafka配置参数
        """

        self.bootstrap_servers = bootstrap_servers
        self.username = kwargs['sasl_plain_username']
        self.password = kwargs['sasl_plain_password']
        self.producer = None
        if self.username and self.password:
            self.producer = KafkaProducer(
                bootstrap_servers=bootstrap_servers,
                value_serializer=lambda v: json.dumps(v, ensure_ascii=False).encode('utf-8'),
                key_serializer=lambda k: k.encode('utf-8') if k else None,
                # 其他常用配置
                acks='all',  # 等待所有副本确认
                retries=3,  # 重试次数
                batch_size=16384,  # 批量发送大小
                linger_ms=10,  # 等待时间
                buffer_memory=33554432,  # 缓冲区大小
                security_protocol='SASL_PLAINTEXT',
                sasl_mechanism='PLAIN',
                sasl_plain_username=self.username,
                sasl_plain_password=self.password,
            )
        else:
            self.producer = KafkaProducer(
                bootstrap_servers=bootstrap_servers,
                value_serializer=lambda v: json.dumps(v, ensure_ascii=False).encode('utf-8'),
                key_serializer=lambda k: k.encode('utf-8') if k else None,
                # 其他常用配置
                acks='all',  # 等待所有副本确认
                retries=3,  # 重试次数
                batch_size=16384,  # 批量发送大小
                linger_ms=10,  # 等待时间
                buffer_memory=33554432,  # 缓冲区大小
                **kwargs
            )

    def validate_kafka_auth(self):
        """使用AdminClient快速验证认证信息"""
        try:
            admin_client = KafkaAdminClient(
                bootstrap_servers=self.bootstrap_servers,
                security_protocol='SASL_PLAINTEXT',
                sasl_mechanism='PLAIN',
                sasl_plain_username=self.username,
                sasl_plain_password=self.password,
                request_timeout_ms=3000
            )

            # 尝试列出topics，这会触发认证
            topics = admin_client.list_consumer_groups()
            admin_client.close()
            print("Kafka认证成功！")
            return True

        except Exception as e:
            print(f"Kafka认证失败: {e}")
            return False

    def send_sync(self, topic, message, key=None, partition=None, timeout=10):
        """
        同步发送消息

        Args:
            topic: 主题名称
            message: 消息内容
            key: 消息键
            partition: 指定分区
            timeout: 超时时间

        Returns:
            RecordMetadata: 发送结果元数据
        """
        try:
            future = self.producer.send(
                topic=topic,
                value=message,
                key=key,
                partition=partition
            )
            record_metadata = future.get(timeout=timeout)
            logging.info(f"消息发送成功 - Topic: {record_metadata.topic}, "
                         f"Partition: {record_metadata.partition}, "
                         f"Offset: {record_metadata.offset}")
            return record_metadata
        except KafkaError as e:
            logging.error(f"消息发送失败: {e}")
            raise

    def send_async(self, topic, message, key=None, partition=None, callback=None):
        """
        异步发送消息

        Args:
            topic: 主题名称
            message: 消息内容
            key: 消息键
            partition: 指定分区
            callback: 回调函数
        """
        try:
            future = self.producer.send(
                topic=topic,
                value=message,
                key=key,
                partition=partition
            )
            if callback:
                future.add_callback(callback)
            future.add_errback(self._error_callback)
        except Exception as e:
            logging.error(f"异步发送消息失败: {e}")

    def _error_callback(self, exc):
        """错误回调函数"""
        logging.error(f"消息发送错误: {exc}")

    def send_batch(self, topic, messages):
        """
        批量发送消息

        Args:
            topic: 主题名称
            messages: 消息列表
        """
        for i, message in enumerate(messages):
            key = f"batch_{i}"
            self.send_async(topic, message, key=key)

        # 等待所有消息发送完成
        self.producer.flush()
        logging.info(f"批量发送完成，共发送 {len(messages)} 条消息")

    def close(self):
        """关闭生产者"""
        self.producer.close()


# 使用示例
def main():
    # 配置日志
    logging.basicConfig(level=logging.INFO)

    # 创建生产者实例
    kafka_producer = KafkaMessageProducer(
        bootstrap_servers=['localhost:9092']
    )

    topic = 'user-events'

    try:
        # 1. 同步发送单条消息
        message = {
            'user_id': 'user_001',
            'event': 'login',
            'timestamp': int(time.time()),
            'ip': '192.168.1.100'
        }

        print("=== 同步发送消息 ===")
        kafka_producer.send_sync(topic, message, key='user_001')

        # 2. 异步发送消息
        print("\n=== 异步发送消息 ===")

        def success_callback(record_metadata):
            print(f"异步发送成功: Topic={record_metadata.topic}, "
                  f"Partition={record_metadata.partition}, "
                  f"Offset={record_metadata.offset}")

        message2 = {
            'user_id': 'user_002',
            'event': 'logout',
            'timestamp': int(time.time()),
            'ip': '192.168.1.101'
        }

        kafka_producer.send_async(topic, message2, key='user_002', callback=success_callback)

        # 3. 批量发送消息
        print("\n=== 批量发送消息 ===")
        batch_messages = []
        for i in range(5):
            msg = {
                'user_id': f'user_{i:03d}',
                'event': 'page_view',
                'page': f'/page_{i}',
                'timestamp': int(time.time())
            }
            batch_messages.append(msg)

        kafka_producer.send_batch(topic, batch_messages)

        # 等待一段时间确保异步消息发送完成
        time.sleep(2)

    except Exception as e:
        print(f"发送消息时出错: {e}")
    finally:
        kafka_producer.close()


if __name__ == "__main__":
    main()
